#include <vector>
#include <numeric>
#include <iostream>
#include <algorithm>
#include <functional>

using namespace std;

// 好路径的数目
// 给你一棵 n 个节点的树（连通无向无环的图）
// 节点编号从0到n-1，且恰好有n-1条边
// 给你一个长度为 n 下标从 0 开始的整数数组 vals
// 分别表示每个节点的值。同时给你一个二维整数数组 edges
// 其中 edges[i] = [ai, bi] 表示节点 ai 和 bi 之间有一条 无向 边
// 好路径需要满足以下条件：开始和结束节点的值相同、 路径中所有值都小于等于开始的值
// 请你返回不同好路径的数目
// 注意，一条路径和它反向的路径算作 同一 路径
// 比方说， 0 -> 1 与 1 -> 0 视为同一条路径。单个节点也视为一条合法路径
// 测试链接 : https://leetcode.cn/problems/number-of-good-paths/

class Solution 
{
private:
    static const int MAXN = 30001;
    // 需要保证集合中，代表节点的值，一定是整个集合的最大值
    int father[MAXN];
    // 集合中最大值的次数，也就是 集合中代表节点的值有几个
    int maxCnt[MAXN];

    void Build(int n)
    {
        for(int i = 0; i < n; ++i)
        {
            father[i] = i;
            maxCnt[i] = 1;
        }
    }

    // 这个并查集的优化只来自扁平化（路径压缩）
    int Find(int i)
    {
        return i == father[i] ? i : father[i] = Find(father[i]);
    }

	// 核心！
	// 注意以下的写法！
	// 谁的值大，谁做代表节点
	// 同时注意 maxCnt 的更新
    int Union(int x, int y, vector<int>& vals)
    {
        // fx : x 所在集合的代表节点，同时也是 x 所在集合的最大值下标
        int fx = Find(x);
        // fy : y 所在集合的代表节点，同时也是 y 所在集合的最大值下标
        int fy = Find(y);
        int path = 0;
        if(vals[fx] > vals[fy])
        {
            father[fy] = fx;
        }
        else if(vals[fx] < vals[fy])
        {
            father[fx] = fy;
        }
        else
        {
            // 两个集合最大值一样！
            path = maxCnt[fx] * maxCnt[fy];
            father[fy] = fx;
            maxCnt[fx] += maxCnt[fy];
        }
        return path;
    }


public:
    int numberOfGoodPaths(vector<int>& vals, vector<vector<int>>& edges) 
    {
        int n = vals.size();
        Build(n);
        // 课上重点讲这个核心排序！
		// 处理边的时候，依次从小节点往大节点处理
        sort(edges.begin(), edges.end(), [&](auto& a, auto& b){
            return max(vals[a[0]], vals[a[1]]) < max(vals[b[0]], vals[b[1]]);
        });

        int ans = n;
        for(auto& e : edges)
        {
            ans += Union(e[0], e[1], vals);
        }
        return ans;
    }
};

// 建图的方式
class Solution2
{
public:
    int numberOfGoodPaths(vector<int>& vals, vector<vector<int>>& edges) 
    {
        int n = vals.size();
        vector<vector<int>> g(n);
        for (auto &e : edges) 
        {
            int x = e[0], y = e[1];
            g[x].push_back(y);
            g[y].push_back(x); // 建图
        }

        // 并查集模板
        // size[x] 表示节点值等于 vals[x] 的节点个数，
        // 如果按照节点值从小到大合并，size[x] 也是连通块内的等于最大节点值的节点个数
        int id[n], fa[n], size[n]; // id 后面排序用
        iota(id, id + n, 0);
        iota(fa, fa + n, 0);
        fill(size, size + n, 1);
        function<int(int)> find = [&](int x) -> int { return fa[x] == x ? x : fa[x] = find(fa[x]); };

        int ans = n; // 单个节点的好路径
        // 根据 val 进行排序
        sort(id, id + n, [&](int i, int j) { return vals[i] < vals[j]; });
        for (int x : id) 
        {
            int vx = vals[x], fx = find(x);
            for (int y : g[x]) // 遍历与 x 相连的节点
            {
                y = find(y);
                if (y == fx || vals[y] > vx)
                    continue; // 只考虑最大节点值不超过 vx 的连通块
                if (vals[y] == vx) // 可以构成好路径
                { 
                    ans += size[fx] * size[y]; // 乘法原理
                    size[fx] += size[y]; // 统计连通块内节点值等于 vx 的节点个数
                }
                fa[y] = fx; // 把小的节点值合并到大的节点值上
            }
        }
        return ans;
    }
};

int main()
{
    Solution obj1;
    Solution2 obj2;

    // 课上例子 1
    vector<int> vals1 = { 2, 1, 1, 2, 2, 1, 1, 2 };
    vector<vector<int>> edges1 = { 
				{ 0, 1 },
				{ 0, 2 },
				{ 1, 3 },
				{ 2, 4 },
				{ 2, 5 },
				{ 5, 6 },
				{ 6, 7 } };
    int ans1 = obj1.numberOfGoodPaths(vals1, edges1);
    int ans2 = obj2.numberOfGoodPaths(vals1, edges1);
    if(ans1 != ans2) cout << "出错了！" << endl;
    else cout << ans1 << endl;

    // 课上例子 2
    vector<int> vals2 = { 1, 2, 2, 3, 1, 2, 2, 1, 1, 3, 3, 3, 3 };
    vector<vector<int>> edges2 = {
				{ 0, 1 },
				{ 0, 2 },
				{ 0, 3 },
				{ 1, 4 },
				{ 4, 7 },
				{ 4, 8 },
				{ 3, 5 },
				{ 3, 6 },
				{ 6, 9 },
				{ 6, 10 },
				{ 6, 11 },
				{ 9, 12 } };
    ans1 = obj1.numberOfGoodPaths(vals2, edges2);
    ans2 = obj1.numberOfGoodPaths(vals2, edges2);
    if(ans1 != ans2) cout << "出错了！" << endl;
    else cout << ans1 << endl;

    return 0;
}
				
				
				
				
				
				